//
// Created by ZhaoXiaoFei on 2022/10/25.
// Last changed by: ZhaoXiaoFei.
// Last update: 2022/10/25.
//

#ifndef MSCKF_MONO_CORNER_DETECTOR_H
#define MSCKF_MONO_CORNER_DETECTOR_H

#include <opencv2/opencv.hpp>
#include <Eigen/Dense>
#include <opencv2/core/eigen.hpp>
#include <random>
#include "data_types.h"

namespace corner_detector{

    class CornerDetector{
    public:
        CornerDetector() = default;
        ~CornerDetector() = default;

        void SetParameters(int grid_rows = 10, int grid_cols = 10, double detection_threshold = 40.0){
            set_parameters_ = true;
            grid_rows_ = grid_rows;
            grid_cols_ = grid_cols;
            detection_threshold_ = detection_threshold;
            grid_flag_.resize(grid_rows_ * grid_cols_, false);
        }

        void SetGridFlag(const cv::Point2f& p){
            grid_flag_[Point2Index(p)] = true;
        }

        int Point2Index(const cv::Point2f& point){
            return static_cast<int>(point.y / grid_height_) * grid_cols_ + static_cast<int>(point.x / grid_width_);
        }

        void DetectFeatures(const cv::Mat& image, std::vector<cv::Point2f>& features){
            if(!set_parameters_){
                std::cout << "Please set parameters firstly!" << std::endl;
                return;
            }
            grid_height_ = (image.rows / grid_rows_) + 1;
            grid_width_ = (image.cols / grid_cols_) + 1;
            features.clear();

            std::vector<cv::KeyPoint> fast_corners;
            cv::Ptr<cv::FeatureDetector> fast = cv::FastFeatureDetector::create(20, true, cv::FastFeatureDetector::TYPE_9_16);
            fast->detect(image, fast_corners);

            std::vector<double> scores_table;
            std::vector<cv::Point2f> temp_points;
            scores_table.resize(grid_rows_ * grid_cols_, 0.0);
            temp_points.resize(grid_rows_ * grid_cols_);

            for(int i = 0; i < fast_corners.size(); i++){
                if(fast_corners[i].pt.x < 0 || fast_corners[i].pt.x >= image.cols
                   || fast_corners[i].pt.y < 0 || fast_corners[i].pt.y >= image.rows) continue;

                const int index = Point2Index(cv::Point2f(fast_corners[i].pt.x, fast_corners[i].pt.y));
                if(grid_flag_[index]){
                    continue;
                }
                const float score = shiTomasiScore(image, fast_corners[i].pt.x, fast_corners[i].pt.y);
                if(score > scores_table[index]){
                    scores_table[index] = static_cast<double>(score);
                    temp_points[index] = cv::Point2f(fast_corners[i].pt.x, fast_corners[i].pt.y);
                }
            }
            for(int i = 0; i < scores_table.size(); i++){
                if(scores_table[i] > detection_threshold_){
                    features.push_back(temp_points[i]);
                }
            }
            ZeroGrid();
        }

        float shiTomasiScore(const cv::Mat &image, int u, int v) {

            float dXX = 0, dYY = 0, dXY = 0;
            const int half_window_size = 15;
            const int window_size = 2 * half_window_size;
            const int window_area = window_size * window_size;
            const float x_min = u - half_window_size;
            const float x_max = u + half_window_size;
            const float y_min = v - half_window_size;
            const float y_max = v + half_window_size;

            if (x_min < 1 || x_max >= image.cols - 1 || y_min < 1 || y_max >= image.rows - 1)
                return 0.0;

            for (float x = x_min; x <= x_max; x++) {
                for (float y = y_min; y <= y_max; y++) {
                    const uint8_t left = image.at<uint8_t>(y, x - 1);
                    const uint8_t right = image.at<uint8_t>(y, x + 1);
                    const uint8_t up = image.at<uint8_t>(y - 1, x);
                    const uint8_t down = image.at<uint8_t>(y + 1, x);
                    float dx = right - left;
                    float dy = down - up;
                    dXX += dx * dx;
                    dYY += dy * dy;
                    dXY += dx * dy;
                }
            }
            dXX /= 2 * window_area;
            dYY /= 2 * window_area;
            dXY /= 2 * window_area;
            // (-b +- sqrt(b * b - 4 * a * c))/ 2 * a;
            return 0.5 * (dXX + dYY - sqrt((dXX + dYY) * (dXX + dYY) - 4 * (dXX * dYY - dXY * dXY)));
        }

        int GetRows(){
            return grid_rows_;
        }
        int GetCols(){
            return grid_cols_;
        }

    private:
        void ZeroGrid(){
            std::fill(grid_flag_.begin(), grid_flag_.end() - 1, false);
        }

    private:
        bool set_parameters_ = false;
        int grid_rows_, grid_cols_, grid_height_, grid_width_;
        std::vector<bool> grid_flag_;
        double detection_threshold_;

    };

    class CornerTracker{
    public:
        CornerTracker(int window_size=51,
                      double min_eigen_threshold=0.001,
                      int max_level=5,
                      int termcrit_max_iters=50,
                      double termcirt_epsilon=0.01)
                : window_size_(cv::Size(window_size, window_size)), max_level_(max_level),
                  min_eigen_threshold_(min_eigen_threshold),
                  termination_criteria_(cv::TermCriteria(cv::TermCriteria::COUNT+cv::TermCriteria::EPS, termcrit_max_iters, termcirt_epsilon)){

        }
        ~CornerTracker() = default;

        void SetParameters(int window_size,
                           double min_eigen_threshold,
                           int max_level,
                           int termcrit_max_iters,
                           double termcirt_epsilon){

            window_size_ = cv::Size(window_size, window_size);
            min_eigen_threshold_ = min_eigen_threshold;
            max_level_ = max_level;
            termination_criteria_ = cv::TermCriteria(cv::TermCriteria::COUNT + cv::TermCriteria::EPS, termcrit_max_iters,
                                                     termcirt_epsilon);

        }

        void TrackFeatures(const cv::Mat& img1, const cv::Mat& img2,
                           std::vector<cv::Point2f>& points1, std::vector<cv::Point2f>& points2,
                           std::vector<size_t>& idx1, std::vector<size_t>& idx2){
            std::vector<uchar> status;
            std::vector<float> err;
            cv::calcOpticalFlowPyrLK(img1, img2, points1, points2, status, err,
                                     window_size_, max_level_, termination_criteria_,
                                     cv::OPTFLOW_USE_INITIAL_FLOW,min_eigen_threshold_);
            for(int i = 0; i < points2.size(); i++){
                if(status[i] && (!inBorder(img2, points2[i]) || distanceTooFar(points1[i], points2[i]))){
                    status[i] = 0;
                }
            }
            reduceVector(points1, status);
            reduceVector(points2, status);
            reduceVector(idx1, status);
            reduceVector(idx2, status);
        }

    private:
        template<class T>
        void reduceVector(std::vector<T> &vec, std::vector<unsigned char> &status){
            int j = 0;
            for(int i = 0; i < vec.size(); i++){
                if(status[i]){
                    vec[j++] = vec[i];
                }
            }
            vec.resize(j);
        }

        bool inBorder(const cv::Mat& image, const cv::Point2f &pt)
        {
            const int BORDER_SIZE = 1;
            int img_x = cvRound(pt.x);
            int img_y = cvRound(pt.y);
            return BORDER_SIZE <= img_x && img_x < image.cols - BORDER_SIZE && BORDER_SIZE <= img_y && img_y < image.rows - BORDER_SIZE;
        }

        bool distanceTooFar(const cv::Point2f &pt1, const cv::Point2f &pt2)
        {
            cv::Point2f dist_vector = pt2 - pt1;
            double dist = std::sqrt(dist_vector.x * dist_vector.x + dist_vector.y * dist_vector.y);
            return dist > 25.0;
        }

    private:
        cv::Size window_size_;
        int max_level_;
        cv::TermCriteria termination_criteria_;
        double min_eigen_threshold_;

    };

    class CornerVisualizer{
    public:
        CornerVisualizer() {
            tracked_points_.clear();
        };
        ~CornerVisualizer() = default;

        void AddPredictedPoints(const std::vector<cv::Point2f>& points, const std::vector<size_t>& idxs){
            assert(points.size() == idxs.size());
            predicted_points_.clear();

            for(int i = 0; i < points.size(); i++){
                predicted_points_.insert(std::make_pair(idxs[i], points[i]));
            }
        }

        void AddNewPoints(const std::vector<cv::Point2f>& points, const std::vector<size_t>& idxs){
            assert(points.size() == idxs.size());
            for(int i = 0; i < points.size(); i++){
                if(tracked_points_.find(idxs[i]) == tracked_points_.end()){
                    tracked_points_.insert(std::make_pair(idxs[i], std::vector<cv::Point2f>()));
                    tracked_points_[idxs[i]].push_back(points[i]);
                }
            }
        }

        void AddTrackedPoints(const std::vector<cv::Point2f>& points, const std::vector<size_t>& idxs){
            assert(points.size() == idxs.size());
            for(int i = 0; i < points.size(); i++){
                if(tracked_points_.find(idxs[i]) == tracked_points_.end()){
                    tracked_points_.insert(std::make_pair(idxs[i], std::vector<cv::Point2f>() ));
                }
                tracked_points_[idxs[i]].push_back(points[i]);
            }
            std::vector<size_t> need_to_remove;
            for(auto& iter : tracked_points_){
                if(std::find(idxs.begin(), idxs.end(), iter.first) == idxs.end()){
                    need_to_remove.push_back(iter.first);
                }
            }
            for(int i = 0; i < need_to_remove.size(); i++){
                tracked_points_.erase(need_to_remove[i]);
            }
        }

        cv::Mat Draw(cv::Mat image){
            cv::Mat res_image;
            if(image.type() == CV_8UC3){
                res_image = image.clone();
            } else if(image.type() == CV_8UC1){
                cv::cvtColor(image, res_image, cv::COLOR_GRAY2BGR);
            }
            else{
                CV_Error(cv::Error::StsBadArg, "Incorrect type of input image.\n");
                return res_image;
            }
            bool isFirst = false;
            cv::Point2f prev_point;
            for(auto& track : tracked_points_){
                isFirst = true;
                cv::Scalar color = cv::Scalar((track.first / 64 % 8) * 255 / 8, (track.first / 8 % 8) * 255 / 8, (track.first % 8) * 255 / 8);

                if(predicted_points_.find(track.first) != predicted_points_.end()){
                    cv::circle(res_image, predicted_points_[track.first], 6, color, 2);
                }
                for(auto it = track.second.rbegin(); it != track.second.rend(); it++){
                    if(isFirst){
                        isFirst = false;
                        cv::circle(res_image, *it, 3, color, 2);
                    }
                    else{
                        cv::line(res_image, *it, prev_point, color, 1);
                    }
                    prev_point = *it;
                }
            }

            return res_image;
        }

    private:
        std::unordered_map<int, cv::Point2f> predicted_points_;
        std::unordered_map<int, std::vector<cv::Point2f>> tracked_points_;
    };

    class CornerTrackHandler{
    public:
        CornerTrackHandler(const cv::Mat K, const cv::Mat distortion_coeffs, int grid_rows, int grid_cols)
                        : K_(K), K_inv_(K.inv()), distortion_coeffs_(distortion_coeffs)
        {
            clear_tracks();
            corner_detector_.SetParameters(grid_rows, grid_cols, 40);
            corner_track_.SetParameters(51, 0.00001, 4, 30, 1.0);
        }
        ~CornerTrackHandler() = default;

        void SetCurrentImage(cv::Mat cur_image, double time){
            pre_time_ = cur_time_;
            pre_image_ = cur_image_;
            pre_features_ = cur_features_;
            pre_ids_ = cur_ids_;

            std::copy(new_features_.begin(), new_features_.end(), std::back_inserter(pre_features_));
            std::copy(new_ids_.begin(), new_ids_.end(), std::back_inserter(pre_ids_));

            //mask valid.
            //NOTE: I think it's OK not to do it
            //assert(pre_features_.size() == pre_idxs_.size());
            std::vector<bool> grid(corner_detector_.GetRows() * corner_detector_.GetCols(), false);
            auto it1 = pre_features_.begin();
            auto it2 = pre_ids_.begin();
            for(; it1 != pre_features_.end() && it2 != pre_ids_.end(); ){
                size_t index = corner_detector_.Point2Index(*it1);
                if(grid[index]){
                    it1 = pre_features_.erase(it1);
                    it2 = pre_ids_.erase(it2);
                }
                else{
                    grid[index] = true;
                    it1++;
                    it2++;
                }
            }

            cur_image_ = cur_image;
            cur_time_ = time;
            cur_features_.clear();
            cur_ids_.clear();
            new_features_.clear();
            new_ids_.clear();

        }

        void TrackFeatures(std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>>& output_features,
                           std::vector<size_t>& output_ids){
            if(pre_features_.size()){
                predict_features();
                corner_visualizer_.AddPredictedPoints(cur_features_, cur_ids_);
                corner_track_.TrackFeatures(pre_image_, cur_image_,
                                           pre_features_, cur_features_,
                                           pre_ids_, cur_ids_);
            }

            corner_visualizer_.AddTrackedPoints(cur_features_, cur_ids_);

            output_features.clear();
            output_ids.clear();
            if(cur_features_.size() > 0){
                std::vector<cv::Point2f> undistorted_pre_features;
                undistortedPoints(pre_features_, undistorted_pre_features);

                std::vector<cv::Point2f> undistorted_cur_features;
                undistortedPoints(cur_features_, undistorted_cur_features);

                std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>> pre_features;
                std::transform(undistorted_pre_features.begin(), undistorted_pre_features.end(), std::back_inserter(pre_features),
                               [](const cv::Point2f &pt) { return Eigen::Vector2f{pt.x, pt.y}; });

                std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>> cur_features;
                std::transform(undistorted_cur_features.begin(), undistorted_cur_features.end(), std::back_inserter(cur_features),
                               [](const cv::Point2f &pt) { return Eigen::Vector2f{pt.x, pt.y}; });


                //RANSAC
                Eigen::Matrix3f dR;
                cv::cv2eigen(dR_, dR);
                if (cur_features.size() > 5 && false) {
                    auto valid = TwoPointsRansac(dR, pre_features, cur_features);

                    auto iter = cur_features.begin();
                    auto cur_features_iter_ = cur_features_.begin();
                    auto cur_ids_iter_ = cur_ids_.begin();

                    for (int i = 0; i < cur_features.size(); i++) {
                        if (!valid[i]) {
                            iter = cur_features.erase(iter);
                            cur_features_iter_ = cur_features_.erase(cur_features_iter_);
                            cur_ids_iter_ = cur_ids_.erase(cur_ids_iter_);
                        }
                        else {
                            iter++;
                            cur_features_iter_++;
                            cur_ids_iter_++;
                        }
                    }
                }

                std::copy(cur_features.begin(), cur_features.end(),
                          std::back_inserter(output_features));
                std::copy(cur_ids_.begin(), cur_ids_.end(),
                          std::back_inserter(output_ids));

            }

        }

        void NewFeatures(std::vector<Eigen::Vector2f,Eigen::aligned_allocator<Eigen::Vector2f>>& output_features,
                         std::vector<size_t>& output_ids){
            //set mask
            for(int i = 0; i < cur_features_.size(); i++){
                corner_detector_.SetGridFlag(cur_features_[i]);
            }

            corner_detector_.DetectFeatures(cur_image_, new_features_);

            new_ids_.clear();
            for(int i = 0; i < new_features_.size(); i++){
                new_ids_.push_back(++global_feature_id);
            }

            corner_visualizer_.AddNewPoints(new_features_, new_ids_);

            output_features.clear();
            output_ids.clear();
            if(new_features_.size()){
                std::vector<cv::Point2f> undisPoints;
                undistortedPoints(new_features_, undisPoints);

                std::transform(undisPoints.begin(), undisPoints.end(), std::back_inserter(output_features),
                               [](const cv::Point2f& p){
                    return Eigen::Vector2f(p.x, p.y);
                });
                std::copy(new_ids_.begin(), new_ids_.end(), std::back_inserter(output_ids));
            }
        }

        void AddGyrData(const Eigen::Vector3f& gyr_data){
            gyr_sum_ += gyr_data;
            gyr_num++;
        }
        void SetRansacParameters(double in){
            ransac_threshold_ = in;
        }

        cv::Mat GetTrackedImage(){
            return corner_visualizer_.Draw(cur_image_);
        }

    private:
        //Node: Returns the coordinates of the normalized plane
        void undistortedPoints(const std::vector<cv::Point2f>& in, std::vector<cv::Point2f>& out){
            cv::undistortPoints(in, out, K_, distortion_coeffs_);
        }
        void clear_tracks(){
            pre_image_ = cv::Mat();
            pre_features_.clear();
            pre_ids_.clear();
        }

        void predict_features() {
            std::copy(pre_ids_.begin(),
                      pre_ids_.end(),
                      std::back_inserter(cur_ids_));
            if(use_gyr){
                integrate_gyro();
                cv::Mat H = K_ * dR_ * K_inv_;
                cv::Mat pt1 = cv::Mat(3, 1, CV_32F);
                cv::Mat pt2 = cv::Mat(3, 1, CV_32F);
                for(int i = 0; i < pre_features_.size(); i++){
                    pt1.at<float>(0, 0) = pre_features_[i].x;
                    pt1.at<float>(1, 0) = pre_features_[i].y;
                    pt1.at<float>(2, 0) = 1;

                    pt2 = H * pt1;

                    cur_features_.emplace_back(pt2.at<float>(0, 0)/pt2.at<float>(2, 0),
                                               pt2.at<float>(1, 0)/pt2.at<float>(2, 0));
                }

            }
            else{
                std::copy(pre_features_.begin(),
                          pre_features_.end(),
                          std::back_inserter(cur_features_));
            }
        }

        void integrate_gyro(){
            double delta_T = cur_time_ - pre_time_;
            if(gyr_num > 0){
                gyr_sum_ /= gyr_num;
            } else{
                use_gyr = false;
                return;
            }
            gyr_sum_ *= delta_T;
            cv::Mat r = cv::Mat(3, 1, CV_32F);
            r.at<float>(0, 0) = gyr_sum_(0, 0);
            r.at<float>(1, 0) = gyr_sum_(1, 0);
            r.at<float>(2, 0) = gyr_sum_(2, 0);

            cv::Rodrigues(r, dR_);

            gyr_sum_.setZero();
            gyr_num = 0;
        }
        Eigen::Array<bool, 1, Eigen::Dynamic>
        TwoPointsRansac(const Eigen::Matrix3f& dR,
                         const std::vector<Eigen::Vector2f, Eigen::aligned_allocator<Eigen::Vector2f>>& pre_points,
                         const std::vector<Eigen::Vector2f, Eigen::aligned_allocator<Eigen::Vector2f>>& cur_points){
            assert(pre_points.size() == cur_points.size());

            int num_points = pre_points.size();
            Eigen::Matrix<float, 3, Eigen::Dynamic> old_points(3, num_points), new_points(3, num_points);
            old_points.setZero(); new_points.setZero();

            for(int i = 0; i < num_points; i++){
                old_points.col(i).x() = pre_points[i].x();
                old_points.col(i).y() = pre_points[i].y();
                old_points.col(i).z() = 1.0;

                new_points.col(i).x() = cur_points[i].x();
                new_points.col(i).y() = cur_points[i].y();
                new_points.col(i).z() = 1.0;
            }

            int iterations = 300;
            std::default_random_engine gen(time(nullptr));
            std::uniform_int_distribution<int> uid(0, num_points - 1);

            int best_num_inliners = -1;
            Eigen::Array<bool, 1, Eigen::Dynamic> best_inliners;

            for(int i = 0; i < iterations; i++){
                int id1 = uid(gen);
                int id2 = id1;
                while(id1 == id2){
                    id2 = uid(gen);
                }
                Eigen::Matrix<float, 3, 1> t;
                Eigen::Matrix<float, 2, 3> M;
                M << new_points.col(id1).transpose() * vectorToSkewSymmetric(dR * old_points.col(id1)),
                        new_points.col(id2).transpose() * vectorToSkewSymmetric(dR * old_points.col(id2));

                if(!M.isZero(1e-9)){
//                    Eigen::FullPivLU<Eigen::Matrix<float, 2, 3>> lu(M);
//                    t = lu.kernel();
                    Eigen::JacobiSVD<Eigen::Matrix<float, 2, 3>> svd(M, Eigen::ComputeFullU | Eigen::ComputeFullV);
                    t = svd.matrixV().rightCols<1>();
                    t /= t(2, 0);
                }else{
                    t.setZero();
                }

                if (t.cols() > 1) {
                    printf("Kernel in RANSAC is the wrong size, returning.");
                    continue;
                }
                Eigen::Matrix3f E = vectorToSkewSymmetric(t) * dR;
                Eigen::Array<float, 3, Eigen::Dynamic> EX1 = E * old_points;
                Eigen::Array<float, 3, Eigen::Dynamic> EX2 = E.transpose() * new_points;

                Eigen::Array<float, Eigen::Dynamic, Eigen::Dynamic> errs = (new_points.array() * EX1).colwise().sum().array().square();

                errs /= EX1.row(0).array().square() + EX1.row(1).array().square() + EX2.row(0).array().square() + EX2.row(1).array().square();
                Eigen::Array<bool, 1, Eigen::Dynamic> inliners = errs < ransac_threshold_;
                int num_inliners = inliners.count();
                if(num_inliners > best_num_inliners){
                    best_num_inliners = num_inliners;
                    best_inliners = inliners;
                }
            }
            std::cout << "best_num_inliners: " << best_num_inliners << std::endl;
            std::cout << "best_inliners.size(): " << best_inliners.size() << std::endl;
            return best_inliners;
        }



    private:
        size_t global_feature_id = 0;
        cv::Mat K_, K_inv_, distortion_coeffs_;
        bool use_gyr = true;
        Eigen::Vector3f gyr_sum_ = Eigen::Vector3f::Zero();
        int gyr_num = 0;
        cv::Mat dR_ = cv::Mat::eye(3, 3, CV_32F);

        CornerDetector corner_detector_;
        CornerTracker corner_track_;
        CornerVisualizer corner_visualizer_;

        cv::Mat cur_image_;
        double cur_time_;
        cv::Mat pre_image_;
        double pre_time_;

        std::vector<cv::Point2f> cur_features_;
        std::vector<size_t> cur_ids_;

        std::vector<cv::Point2f> new_features_;
        std::vector<size_t> new_ids_;

        std::vector<cv::Point2f> pre_features_;
        std::vector<size_t> pre_ids_;

        double ransac_threshold_;

    };

}


#endif //MSCKF_MONO_CORNER_DETECTOR_H
